import numpy as np
import cv2
from facility.pic_fn import Pic,rgb_to_binary


def yanzhengma(counters):
    res = np.array([i for i in counters if 85 < i.shape[0] < 110])
    res = [i for i in res if np.max(i[:,:,1])-np.min(i[:,:,1]) < 6]
    return res



def find_contours(image,filters=None):
    # cv2.CHAIN_APPROX_NONE存储所有的轮廓点，相邻的两个点的像素位置差不超过1，即max（abs（x1-x2），abs（y2-y1））==1
    # cv2.CHAIN_APPROX_SIMPLE压缩水平方向，垂直方向，对角线方向的元素，只保留该方向的终点坐标，例如一个矩形轮廓只需4个点来保存轮廓信息
    origin, counters, rank = cv2.findContours(image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    if filters is not None:
        counters = filters(counters)
    return counters


def display_counters(orignal,counters):

    cot = Pic.gen(orignal.shape)()

    cot[:, :] = True
    for counter in counters:
        for i in counter:
            coor = i.flatten()
            cot[coor[1], coor[0]] = False

    Pic.load(cot).display()
    return cot



def check_mask(frame):
    data =frame
    data[0:20,:,:] =  [0, 0, 0]
    data[545:, :, :] = [0, 0, 0]
    data[:, :70, :] = [0, 0, 0]
    data[484:555, 70:200, :] = [0, 0, 0]
    data[:196, 824:, :] = [0, 0, 0]
    data = rgb_to_binary(data,127)

    res = find_contours(data,yanzhengma)
    if res.__len__() != 6 :
        return None


    points = [np.array(np.average(i, axis=0).flatten(), dtype=np.int32) for i in res]


    part = []
    for point in points:
        x = point[0]
        y = point[1]
        coco = frame[y + 3:y + 14, x - 20:x + 20, :]

        part.append(coco)

    sums = [np.sum(i) for i in part]
    avg = np.average(sums)

    diff = [abs(i - avg) for i in sums]
    loc = np.where(np.max(diff))[0][0]
    return points[loc]


if __name__ == '__main__':
    g = Pic.load('../ma1.png')()
    coor = check_mask(g)
    print(coor)

